#ifndef __PA_MODEL_QWEN2_HUB_H__
#define __PA_MODEL_QWEN2_HUB_H__

#include "model/model_base.h"

#include <string>
#include <nlohmann/json.hpp>

using Json = nlohmann::json;

namespace pa
{

class Qwen2Hub: public ModelBase
{
public:
    Qwen2Hub(const Type model_type, const std::string& hub);
    ~Qwen2Hub();

    bool is_valid() const;
    bool load();

    inline size_t token_size() const
    {
        return m_token_strs.size();
    }

    inline std::string& pre_token_pattern()
    {
        return m_pre_token_pattern;
    }

    inline const std::vector<std::string>& token_str() const
    {
        return m_token_strs;
    }

    inline const std::vector<std::string>& added_token_str() const
    {
        return m_added_token_strs;
    }

private:

    // id to token, decoder
    std::vector<std::string> m_token_strs;
    std::vector<std::string> m_added_token_strs;

    std::map<std::string, std::string> m_special_token_strs = {
        {"eos_token", "<|endoftext|>"},
        {"pad_token", "<|endoftext|>"}
    };

    std::vector<std::string> m_additional_special_tokens = {"<|im_end|>", "<|im_end|>"};

    static std::vector<std::string> m_basic_files;
    Json m_tokenizer_config_json;
    Json m_tokenizer_json;
    Json m_vocab_json;

    std::string m_config_tokenizer_class;
    std::string m_pre_token_pattern;
};



} // namespace pa

#endif // __PA_MODEL_QWEN2_HUB_H__